package com.xs.ai.config;

import com.xs.ai.services.BookingTools;
import com.xs.ai.services.FlightBookingService;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.CommandLineRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Description;
import org.springframework.core.io.Resource;

import java.util.function.Function;

/**
 * @author liaohao
 * @createDate 2025/2/11 15:58
 **/
@Configuration
public class ChatConfig {
    @Autowired
    private FlightBookingService flightBookingService;

    /**
     * 注入chatMemory,用于记录对话
     *
     * @return
     */
    @Bean
    public ChatMemory chatMemory() {
        return new InMemoryChatMemory();
    }

    //Jdk17 新语法
    public record CancelBookingRequest(String bookingNumber, String customerName) {
    }

    public record BookingDetailRequest(String bookingNumber, String customerName) {
    }

    public record changeBookingRequest(String bookingNumber, String customerName, String newDate, String from, String to) {
    }

    /**
     * @Description: 处理退订的逻辑
     * @author liao hao
     * @createDate 2025/2/11 18:06
     **/
    @Bean
    @Description("处理机票退订")//描述-用户输入的指令
    public Function<CancelBookingRequest, String> cancelBookingRequestFunction() {
        //apply方法调用，处理用户输入的指令，调用退订功能
        return cancelBookingRequest -> {
            try {
                flightBookingService.cancelBooking(cancelBookingRequest.bookingNumber(), cancelBookingRequest.customerName());
                return "退订成功";
            }catch (Exception e){
                return "退订失败";
            }
        };
    }

    /**
     * @Description: 处理用户问题，返回订单信息
     * @author liao hao
     * @createDate 2025/2/11 18:06
     **/
    @Bean
    @Description("获取机票预定详细信息")//描述-用户输入的指令
    public Function<BookingDetailRequest, BookingTools.BookingDetails> getBookingRequestFunction() {
        //apply方法调用，处理用户输入的指令，调用订单详细功能
        return getBookingRequest -> {
            try {
                return flightBookingService.getBookingDetails(getBookingRequest.bookingNumber(), getBookingRequest.customerName());
            } catch (Exception e) {
                return new BookingTools.BookingDetails(getBookingRequest.bookingNumber(), getBookingRequest.customerName(), null, null, null, null, null);
            }
        };
    }

    @Bean
    @Description("更新机票预定信息")//描述-用户输入的指令
    public Function<changeBookingRequest, String> updateBookingRequestFunction() {
        return updateBookingRequest -> {
            try {
                flightBookingService.changeBooking(updateBookingRequest.bookingNumber(), updateBookingRequest.customerName(), updateBookingRequest.newDate(), updateBookingRequest.from(), updateBookingRequest.to());
                return "更新成功";
            }catch (Exception e){
                return "更新失败";
            }
        };
    }

    /**
     * RAG
     *
     * @Description: 注入向量数据库
     * @author liao hao
     * @createDate 2025/2/12 11:48
     **/
    @Bean
    public VectorStore vectorStore(EmbeddingModel embeddingModel) {
        return new SimpleVectorStore(embeddingModel);
    }

    @Bean
    public CommandLineRunner ingestTermOfServiceToVectorStore(EmbeddingModel embeddingModel, VectorStore vectorStore,
                                                              @Value("classpath:rag/terms-of-service.txt") Resource termsOfServiceDocs) {
        return args -> {
            vectorStore.write( // 3.写⼊
                    new TokenTextSplitter().transform( // 2.转换
                            new TextReader(termsOfServiceDocs).read()) // 1.读取
            );
        };
    }
}
